package bge

import (
	"context"
	"fmt"
	"github.com/go-resty/resty/v2"
	"knowledge_admin/pkg/framework"
	"time"
)

const (
	defaultTimeout = 10 * time.Second
	rerankEndpoint = "/v1/rerank"
	maxRetries     = 3
	retryDelay     = 500 * time.Millisecond
)

type BgeRerankClient struct {
	baseURL string
	client  *resty.Client
	logger  framework.Logger
}

type BaseRespData struct {
	Code int         `json:"code"`
	Msg  interface{} `json:"message"`
	Data interface{} `json:"data"`
}

func NewBgeRerankClient(baseURL string, logger framework.Logger) *BgeRerankClient {
	client := resty.New().
		SetTimeout(defaultTimeout).
		SetRetryCount(maxRetries).
		SetRetryWaitTime(retryDelay).
		SetHeader("Content-Type", "application/json")

	return &BgeRerankClient{
		baseURL: baseURL,
		client:  client,
		logger:  logger,
	}
}

func (c *BgeRerankClient) Rerank(ctx context.Context, req, resp interface{}) error {
	url := c.baseURL + rerankEndpoint
	ret := &BaseRespData{Data: resp}
	r, err := c.client.R().
		SetContext(ctx).
		SetBody(req).
		SetResult(&ret).
		Post(url)

	if err != nil {
		c.logger.Error("failed to send request", err)
		return fmt.Errorf("failed to send request: %w", err)
	}

	if r.StatusCode() != 200 {
		c.logger.Error("unexpected status code", fmt.Errorf("status code: %d", r.StatusCode()))
		return fmt.Errorf("unexpected status code: %d", r.StatusCode())
	}
	if ret.Code != 0 {
		c.logger.Error("rerank failed", fmt.Errorf("code: %d, message: %v", ret.Code, ret.Msg))
		return fmt.Errorf("rerank failed: %d, %v", ret.Code, ret.Msg)
	}
	return nil
}
